package com.agentsflex;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Scanner;

public class OllamaChat {
    private static final String OLLAMA_API_URL = "http://localhost:11434/api/generate";
    private static final String MODEL_NAME = "deepseek-r1:8b";
    private static final ObjectMapper mapper = new ObjectMapper();

    private final List<String> chatHistory = new ArrayList<>();

    public static void main(String[] args) throws IOException {
        OllamaChat chat = new OllamaChat();
        Scanner scanner = new Scanner(System.in);

        System.out.println("开始对话（输入'exit'退出）:");
        while (true) {
            System.out.print("用户: ");
            String input = scanner.nextLine();

            if ("exit".equalsIgnoreCase(input)) break;

            chat.sendChatRequest(input);
        }
        scanner.close();
    }

    private void sendChatRequest(String userInput) throws IOException {
        // 添加用户输入到历史
        chatHistory.add("用户: " + userInput);

        // 构造请求体
        String prompt = String.join("\n", chatHistory) + "\n助手: ";
        String requestBody = String.format(
            "{\"model\": \"%s\", \"prompt\": \"%s\", \"stream\": true}",
            MODEL_NAME, escapeJson(prompt));

        // 发送HTTP请求
        HttpURLConnection connection = (HttpURLConnection) new URL(OLLAMA_API_URL).openConnection();
        connection.setRequestMethod("POST");
        connection.setRequestProperty("Content-Type", "application/json");
        connection.setDoOutput(true);

        try (OutputStream os = connection.getOutputStream()) {
            os.write(requestBody.getBytes());
            os.flush();
        }
//        System.out.println("打印请求json:"+requestBody);
        // 处理响应
        if (connection.getResponseCode() != 200) {
            throw new IOException("请求失败: HTTP " + connection.getResponseCode());
        }

        StringBuilder responseText = new StringBuilder();
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(connection.getInputStream()))) {
            String line;
            while ((line = reader.readLine()) != null) {
                if (!line.isEmpty()) {
                    JsonNode responseJson = mapper.readTree(line);
                    String partialResponse = responseJson.get("response").asText();
                    responseText.append(partialResponse);
                    System.out.print(partialResponse);
                    System.out.flush();
                }
            }
        }
        System.out.println();

        // 添加助手响应到历史
        chatHistory.add("助手: " + responseText.toString());
    }

    // 处理JSON特殊字符转义
    private String escapeJson(String input) {
        return input.replace("\\", "\\\\")
            .replace("\"", "\\\"")
            .replace("\n", "\\n")
            .replace("\r", "\\r")
            .replace("\t", "\\t");
    }
}
